# -*- coding: utf-8 -*-
"""
端口监控页签
提供端口vs进程的查询UI和功能
"""

# 在任何tkinter导入之前设置环境
import sys
import os
from modules.tkinter_setup import setup_tcltk_environment
setup_tcltk_environment()

import tkinter as tk
from tkinter import ttk, messagebox, scrolledtext
import psutil


class PortMonitorTab:
    """端口监控页签类"""
    
    def __init__(self, parent_notebook, root):
        """
        初始化端口监控页签
        
        Args:
            parent_notebook: 父级Notebook控件
            root: 主窗口根对象
        """
        self.notebook = parent_notebook
        self.root = root
        
        # 创建页签
        self.frame = ttk.Frame(self.notebook)
        self.notebook.add(self.frame, text="端口vs进程")
        
        self._setup_ui()
    
    def _setup_ui(self):
        """设置UI界面"""
        # 顶部搜索栏
        search_frame = ttk.Frame(self.frame)
        search_frame.pack(fill=tk.X, padx=10, pady=10)
        
        # 端口查询
        ttk.Label(search_frame, text="端口号:").pack(side=tk.LEFT, padx=5)
        self.port_entry = ttk.Entry(search_frame, width=15)
        self.port_entry.pack(side=tk.LEFT, padx=5)
        ttk.Button(search_frame, text="查询端口", command=self.search_port).pack(side=tk.LEFT, padx=5)
        
        # 进程名查询
        ttk.Label(search_frame, text="进程名:").pack(side=tk.LEFT, padx=(20, 5))
        self.process_name_entry = ttk.Entry(search_frame, width=15)
        self.process_name_entry.pack(side=tk.LEFT, padx=5)
        ttk.Button(search_frame, text="查询进程", command=self.search_process_ports).pack(side=tk.LEFT, padx=5)
        
        # 其他按钮
        ttk.Button(search_frame, text="查询所有", command=self.search_all_ports).pack(side=tk.LEFT, padx=(20, 5))
        ttk.Button(search_frame, text="清空", command=self.clear_results).pack(side=tk.LEFT, padx=5)
        
        # 统计信息
        self.stats_label = ttk.Label(search_frame, text="")
        self.stats_label.pack(side=tk.LEFT, padx=20)
        
        # 结果列表
        result_frame = ttk.LabelFrame(self.frame, text="端口使用情况")
        result_frame.pack(fill=tk.BOTH, expand=True, padx=10, pady=(0, 10))
        
        # 创建Treeview
        tree_scroll_y = ttk.Scrollbar(result_frame)
        tree_scroll_y.pack(side=tk.RIGHT, fill=tk.Y)
        
        tree_scroll_x = ttk.Scrollbar(result_frame, orient=tk.HORIZONTAL)
        tree_scroll_x.pack(side=tk.BOTTOM, fill=tk.X)
        
        self.tree = ttk.Treeview(result_frame, 
                                 columns=('进程名', 'PID', '协议', '本地地址', '远程地址', '状态'),
                                 show='headings',
                                 yscrollcommand=tree_scroll_y.set,
                                 xscrollcommand=tree_scroll_x.set)
        
        self.tree.heading('进程名', text='进程名')
        self.tree.heading('PID', text='PID')
        self.tree.heading('协议', text='协议')
        self.tree.heading('本地地址', text='本地地址')
        self.tree.heading('远程地址', text='远程地址')
        self.tree.heading('状态', text='状态')
        
        self.tree.column('进程名', width=150)
        self.tree.column('PID', width=80)
        self.tree.column('协议', width=60)
        self.tree.column('本地地址', width=150)
        self.tree.column('远程地址', width=150)
        self.tree.column('状态', width=120)
        
        self.tree.pack(fill=tk.BOTH, expand=True, padx=5, pady=5)
        tree_scroll_y.config(command=self.tree.yview)
        tree_scroll_x.config(command=self.tree.xview)
        
        # 右键菜单
        self.context_menu = tk.Menu(self.tree, tearoff=0)
        self.context_menu.add_command(label="复制进程名", command=lambda: self._copy_info('进程名'))
        self.context_menu.add_command(label="复制PID", command=lambda: self._copy_info('PID'))
        self.context_menu.add_command(label="复制本地地址", command=lambda: self._copy_info('本地地址'))
        self.context_menu.add_separator()
        self.context_menu.add_command(label="结束进程", command=self._kill_process)
        
        self.tree.bind('<Button-3>', self._show_context_menu)
        
        # 说明文字
        info_frame = ttk.Frame(self.frame)
        info_frame.pack(fill=tk.X, padx=10, pady=(0, 10))
        
        info_text = "状态说明: LISTEN(监听), ESTABLISHED(已连接), TIME_WAIT(等待), CLOSE_WAIT(关闭等待)"
        ttk.Label(info_frame, text=info_text, foreground='gray').pack(side=tk.LEFT)
    
    def search_port(self):
        """根据端口号查询进程"""
        port_str = self.port_entry.get().strip()
        if not port_str:
            messagebox.showwarning("警告", "请输入端口号")
            return
        
        try:
            port = int(port_str)
        except ValueError:
            messagebox.showerror("错误", "端口号必须是数字")
            return
        
        self._clear_tree()
        
        count = 0
        try:
            connections = psutil.net_connections(kind='inet')
            for conn in connections:
                if conn.laddr.port == port:
                    self._add_connection_to_tree(conn)
                    count += 1
            
            self.stats_label.config(text=f"找到 {count} 个连接使用端口 {port}")
            
            if count == 0:
                messagebox.showinfo("提示", f"没有找到使用端口 {port} 的进程")
        
        except Exception as e:
            messagebox.showerror("错误", f"查询失败: {e}")
    
    def search_process_ports(self):
        """根据进程名查询端口使用情况"""
        process_name = self.process_name_entry.get().strip()
        if not process_name:
            messagebox.showwarning("警告", "请输入进程名")
            return
        
        self._clear_tree()
        
        try:
            # 查找匹配的进程
            matching_pids = []
            for proc in psutil.process_iter(['name', 'pid']):
                try:
                    if proc.info['name'] and process_name.lower() in proc.info['name'].lower():
                        matching_pids.append(proc.info['pid'])
                except (psutil.NoSuchProcess, psutil.AccessDenied):
                    continue
            
            if not matching_pids:
                messagebox.showinfo("提示", f"没有找到进程名包含 '{process_name}' 的进程")
                self.stats_label.config(text="")
                return
            
            # 查询这些进程使用的端口
            connections = psutil.net_connections(kind='inet')
            count = 0
            port_set = set()
            
            for conn in connections:
                if conn.pid in matching_pids:
                    self._add_connection_to_tree(conn)
                    port_set.add(conn.laddr.port)
                    count += 1
            
            self.stats_label.config(
                text=f"找到 {len(matching_pids)} 个进程, 使用 {len(port_set)} 个端口, 共 {count} 个连接"
            )
            
            if count == 0:
                messagebox.showinfo("提示", f"进程 '{process_name}' 没有使用任何端口")
        
        except Exception as e:
            messagebox.showerror("错误", f"查询失败: {e}")
    
    def search_all_ports(self):
        """查询所有端口使用情况"""
        self._clear_tree()
        
        try:
            connections = psutil.net_connections(kind='inet')
            count = 0
            
            for conn in connections:
                self._add_connection_to_tree(conn)
                count += 1
            
            self.stats_label.config(text=f"共找到 {count} 个网络连接")
        
        except Exception as e:
            messagebox.showerror("错误", f"查询失败: {e}")
    
    def clear_results(self):
        """清空查询结果"""
        self._clear_tree()
        self.stats_label.config(text="")
        self.port_entry.delete(0, tk.END)
        self.process_name_entry.delete(0, tk.END)
    
    def _add_connection_to_tree(self, conn):
        """添加连接信息到树形列表"""
        try:
            # 获取进程信息
            if conn.pid:
                proc = psutil.Process(conn.pid)
                proc_name = proc.name()
            else:
                proc_name = "System"
            
            # 协议类型
            protocol = "TCP" if conn.type == 1 else "UDP"
            
            # 本地地址
            local_addr = f"{conn.laddr.ip}:{conn.laddr.port}"
            
            # 远程地址
            if conn.raddr:
                remote_addr = f"{conn.raddr.ip}:{conn.raddr.port}"
            else:
                remote_addr = "-"
            
            # 状态
            status = conn.status if conn.status else "-"
            
            self.tree.insert('', 'end', values=(
                proc_name,
                conn.pid if conn.pid else "-",
                protocol,
                local_addr,
                remote_addr,
                status
            ))
        except (psutil.NoSuchProcess, psutil.AccessDenied):
            pass
    
    def _clear_tree(self):
        """清空树形列表"""
        for item in self.tree.get_children():
            self.tree.delete(item)
    
    def _show_context_menu(self, event):
        """显示右键菜单"""
        item = self.tree.identify_row(event.y)
        if item:
            self.tree.selection_set(item)
            self.context_menu.post(event.x_root, event.y_root)
    
    def _copy_info(self, column):
        """复制信息"""
        selection = self.tree.selection()
        if not selection:
            return
        
        item = selection[0]
        values = self.tree.item(item, 'values')
        
        column_index = {
            '进程名': 0,
            'PID': 1,
            '协议': 2,
            '本地地址': 3,
            '远程地址': 4,
            '状态': 5
        }
        
        if column in column_index:
            text = values[column_index[column]]
            self.root.clipboard_clear()
            self.root.clipboard_append(text)
            messagebox.showinfo("成功", f"已复制: {text}")
    
    def _kill_process(self):
        """结束进程"""
        selection = self.tree.selection()
        if not selection:
            return
        
        item = selection[0]
        values = self.tree.item(item, 'values')
        pid = values[1]
        proc_name = values[0]
        
        if pid == "-" or pid == "System":
            messagebox.showwarning("警告", "无法结束系统进程")
            return
        
        if messagebox.askyesno("确认", f"确定要结束进程 {proc_name} (PID: {pid}) 吗?"):
            try:
                proc = psutil.Process(int(pid))
                proc.terminate()
                messagebox.showinfo("成功", f"已发送终止信号给进程 {proc_name}")
                # 刷新列表
                port_str = self.port_entry.get().strip()
                if port_str:
                    self.search_port()
                else:
                    self.search_all_ports()
            except Exception as e:
                messagebox.showerror("错误", f"结束进程失败: {e}")
